# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import json
import logging
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.timestamp import get_timestamp
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.gopax.ws.config import group_symbol
from coin.exchange.gopax.kr_rest.product import GopaxProduct


def _find_pair_id(subscribe_symbol, assets_data, trading_pairs_data):
  base, quote = subscribe_symbol.split('-')
  quote_asset_id = None
  base_asset_id = None
  for asset in assets_data:
    if quote == asset['isoAlpha3']:
      quote_asset_id = asset['assetId']
    elif base == asset['isoAlpha3']:
      base_asset_id = asset['assetId']
    if all((quote_asset_id, base_asset_id)):
      break
  for pair in trading_pairs_data:
    if (base_asset_id == pair['baseAssetId'] and quote_asset_id == pair['quoteAssetId']):
      return pair['tradingPairId']
  raise ValueError(subscribe_symbol, assets_data, trading_pairs_data)


class GopaxSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)
    assert len(self._symbol) == 1
    self._symbol = self._symbol[0]

    self._register_url('wss://wss.gopax.co.kr/primus?_primuscb=MnGL-WK')
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('gopax_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('gopax_ws_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'depth': queue.get_writer(topic_id=topic_depth[1]),
        'trades': queue.get_writer(topic_id=topic_trade[1]),
    }

    self._enable_topic_event_publisher(queue, 'gopax_ws_topic_event', [topic_depth[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    yield self._query_assets()
    yield self._query_trading_pairs()
    yield self._send_subscribe_request()

  @gen.coroutine
  def _query_assets(self):
    # Gopax double encode json, crazy !!
    exchange_request_str = json.dumps(json.dumps({"exchange": 1}))
    yield self._ws.write_message(exchange_request_str)
    yield gen.sleep(0.1)
    query_assets_request_str = json.dumps(json.dumps({"m": 0, "i": 4, "n": "GetAssets", "o": "{}"}))
    yield self._ws.write_message(query_assets_request_str)
    # There could be ping pong messages coming
    for _ in range(10):
      resp = yield self._ws.read_message()
      try:
        resp = json.loads(json.loads(resp))
      except json.decoder.JSONDecodeError:
        self._logger.error('Unknown message: %s' % resp)
        continue
      if 'n' in resp and resp['n'] == 'GetAssets':
        # Gopax crazy again!
        o = json.loads(resp['o'])
        assert o['result']
        self._assets_data = o['data']
        self._publisher['depth'].write(get_timestamp(), resp)
        break

  @gen.coroutine
  def _query_trading_pairs(self):
    query_trading_pairs_request_str = json.dumps(
        json.dumps({
            "m": 0, "i": 10, "n": "GetTradingPairs", "o": "{}"
        }))
    yield self._ws.write_message(query_trading_pairs_request_str)
    for _ in range(10):
      resp = yield self._ws.read_message()
      try:
        resp = json.loads(json.loads(resp))
      except json.decoder.JSONDecodeError:
        self._logger.error('Unknown message: %s' % resp)
        continue
      if 'n' in resp and resp['n'] == 'GetTradingPairs':
        o = json.loads(resp['o'])
        assert o['result']
        self._trading_pairs_data = o['data']
        self._publisher['depth'].write(get_timestamp(), resp)
        break

  @gen.coroutine
  def _send_subscribe_request(self):
    pair_id = _find_pair_id(self._symbol, self._assets_data, self._trading_pairs_data)
    self._pair_id = pair_id
    subscribe_depth_str = json.dumps(
        json.dumps({
            "m": 0, "i": 22, "n": "GetOrderBook", "o": json.dumps({"pairId": pair_id})
        }))
    yield self._ws.write_message(subscribe_depth_str)
    yield gen.sleep(0.1)

    start_str = json.dumps(
        json.dumps({
            "m": 0,
            "i": 20,
            "n": "SubscribeChart",
            "o": json.dumps({
                "pairId": pair_id, "interval": 5
            })
        }))
    yield self._ws.write_message(start_str)
    self._publish_topic_event('START')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return

    if isinstance(message, str):
      if message.startswith('"primus::ping::'):
        pong_msg = message.replace('ping', 'pong')
        self._ws.write_message(pong_msg)
      else:
        self._logger.error('Unknown message: %s' % message)
    elif isinstance(message, dict):
      event = message.get('n', None)
      if event is not None:
        if event == 'GetOrderBook':
          message['pairId'] = self._pair_id
          self._publisher['depth'].write(timestamp, message)
        elif event == 'OrderBookEvent':
          self._publisher['depth'].write(timestamp, message)
        elif event == 'TradeEvent':
          self._publisher['trades'].write(timestamp, message)
        elif event in ['ChartDataEvent', 'TickerEvent']:
          self._logger.debug('Ignored message: %s' % message)
        else:
          self._logger.error('Unknown message: %s' % message)
      else:
        self._logger.error('Unknown message: %s' % message)

  def _decode(self, ws_msg):
    try:
      return json.loads(json.loads(ws_msg))
    except json.decoder.JSONDecodeError:
      # Most likely ping message.
      return ws_msg


class GopaxMultiSubscriberWrapper(object):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbols: Union[List[str], str],
               close_callback=None,
               logger=None):
    assert group_id == "TBD"
    self._subscribers = []
    self._symbols = []
    for product in to_list(symbols):
      self._symbols.append(GopaxProduct.FromStr(product).native_subscription_symbol)

    for symbol in self._symbols:
      group_id = None
      for k, v in group_symbol.items():
        if symbol in v:
          group_id = k
      assert group_id is not None, (symbol, group_id)
      self._subscribers.append(
          GopaxSubscriber(queue, group_id, worker_id, symbol, close_callback, logger))

  def subscribe(self):
    for subscriber in self._subscribers:
      subscriber.subscribe()

  def start(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.start(*args, **kwargs)

  def stop(self, *args, **kwargs):
    for subscriber in self._subscribers:
      subscriber.stop(*args, **kwargs)
